package rs_coder

import (
	"common/rs"
	"github.com/vivint/infectious"
)

type rsCoder struct {
	fc *infectious.FEC
}

// NewRSScheme returns a Reed-Solomon-based ErasureScheme.
func NewRSScheme(rs rs.Rs) RsCoder {
	fc, _ := infectious.NewFEC(int(rs.Required), int(rs.Total))
	return &rsCoder{fc: fc}
}

func (s *rsCoder) Encode(input []byte) (*Pieces, error) {
	dataAlign := s.fc.Required()
	ramain := len(input) % dataAlign
	paddingLen := 0
	if ramain > 0 {
		paddingLen = dataAlign - ramain
	}
	r := &Pieces{
		Shared:     make(map[int][]byte),
		PaddingLen: paddingLen,
	}
	if paddingLen > 0 {
		padding := make([]byte, paddingLen)
		for i := range padding {
			padding[i] = 0
		}
		input = append(input, padding...)
	}
	//sumSize := 0
	//inputSize := len(input)
	err := s.fc.Encode(input, func(s infectious.Share) {
		r.Shared[s.Number] = append([]byte(nil), s.Data...)
		//sumSize += len(s.Data)
	})
	//fmt.Println("rs_coder size ", inputSize, sumSize, " rate ", float32(sumSize)/float32(inputSize))
	if err != nil {
		return r, err
	}
	return r, nil
}

func (s *rsCoder) Decode(rs *Pieces) ([]byte, error) {
	in := rs.Shared
	shares := make([]infectious.Share, 0, len(in))
	for num, data := range in {
		if data != nil && len(data) > 0 {
			shares = append(shares, infectious.Share{Number: num, Data: data})
		}
	}
	data, err := s.fc.Decode(nil, shares)
	if err != nil {
		return nil, err
	}
	n := len(data)
	return data[:n-rs.PaddingLen], nil
}

func (s *rsCoder) TotalCount() int {
	return s.fc.Total()
}

func (s *rsCoder) RequiredCount() int {
	return s.fc.Required()
}
